/* 
 *  Copyright (c) 2020-2022 went_forward. All rights reserved.
 *  Licensed under the MIT License.
 *  See LICENSE file for more information.
 */

#include <map>
#include <vector>
#include <string>
#include <sys/stat.h>

#include "dr_api.h"
#include "drmgr.h"
#include "drcctlib.h"

using namespace std;

#define DRCCTLIB_PRINTF(_FORMAT, _ARGS...) \
    DRCCTLIB_PRINTF_TEMPLATE("mem_check", _FORMAT, ##_ARGS)
#define DRCCTLIB_EXIT_PROCESS(_FORMAT, _ARGS...) \
    DRCCTLIB_CLIENT_EXIT_PROCESS_TEMPLATE("mem_check", _FORMAT, ##_ARGS)

#define SAMPLE_RUN
#ifdef SAMPLE_RUN
#    define UNITE_NUM 1000000000
#    define SAMPLE_NUM 100000000
#endif

#define OUTPUT_SIZE 200
#define MAX_CLIENT_CCT_PRINT_DEPTH 10

static std::string g_folder_name;
static int tls_idx;

typedef struct _per_thread_t {
    int pt_id;
    vector<context_handle_t> *access_NULL_Fault_1;
    vector<context_handle_t> *access_unalloc_addr_Fault_2;
    vector<context_handle_t> *access_freed_addr_Fault_3;
    vector<context_handle_t> *access_uninit_addr_Fault_4;
    file_t output_file;
// #define DEBUG_REUSE
#ifdef DEBUG_REUSE
    file_t log_file;
#endif
} per_thread_t;

void
PrintResult(per_thread_t *pt)
{
    dr_fprintf(pt->output_file, "Thread id : %d\n", pt->pt_id);
    dr_fprintf(pt->output_file, "-------------------------------------------------\n");
    /*if (!(pt->access_NULL_Fault_1)->empty()) {
        int fault_1_num = (pt->access_NULL_Fault_1)->size();
        dr_fprintf(pt->output_file, "Fault 1 : %d problems\n", fault_1_num);
        for (int i = 0; i < fault_1_num; i++) {
            dr_fprintf(pt->output_file, "#%d  create handle : \n", i);
            drcctlib_print_backtrace(pt->output_file, (pt->access_NULL_Fault_1)->at(i), true, true, MAX_CLIENT_CCT_PRINT_DEPTH);
            dr_fprintf(pt->output_file, "***************************************************\n\n");
        }
    } else{
        dr_fprintf(pt->output_file, "***No Fault 1***\n\n");
    }*/

    if (!(pt->access_unalloc_addr_Fault_2)->empty()) {
        int fault_2_num = (pt->access_unalloc_addr_Fault_2)->size();
        dr_fprintf(pt->output_file, "Fault 2 : %d problems\n", fault_2_num);
        for (int i = 0; i < fault_2_num; i++) {
            dr_fprintf(pt->output_file, "#%d  create handle : \n", i);
            drcctlib_print_backtrace(pt->output_file, (pt->access_unalloc_addr_Fault_2)->at(i), true, true, MAX_CLIENT_CCT_PRINT_DEPTH);
            dr_fprintf(pt->output_file, "***************************************************\n\n");
        }
    } else {
        dr_fprintf(pt->output_file, "***No Fault 2***\n\n");
    }

    if (!(pt->access_freed_addr_Fault_3)->empty()) {
        int fault_3_num = (pt->access_freed_addr_Fault_3)->size();
        dr_fprintf(pt->output_file, "Fault 3 : %d problems\n", fault_3_num);
        for (int i = 0; i < fault_3_num; i++) {
            dr_fprintf(pt->output_file, "#%d  create handle : \n", i);
            drcctlib_print_backtrace(pt->output_file, (pt->access_freed_addr_Fault_3)->at(i), true, true, MAX_CLIENT_CCT_PRINT_DEPTH);
            dr_fprintf(pt->output_file, "***************************************************\n\n");
        }
    } else {
        dr_fprintf(pt->output_file, "***No Fault 3***\n\n");
    }

    if (!(pt->access_uninit_addr_Fault_4)->empty()) {
        int fault_4_num = (pt->access_uninit_addr_Fault_4)->size();
        dr_fprintf(pt->output_file, "Fault 4 : %d problems\n", fault_4_num);
        for (int i = 0; i < fault_4_num; i++) {
            dr_fprintf(pt->output_file, "#%d  create handle : \n", i);
            drcctlib_print_backtrace(pt->output_file, (pt->access_uninit_addr_Fault_4)->at(i), true, true, MAX_CLIENT_CCT_PRINT_DEPTH);
            dr_fprintf(pt->output_file, "***************************************************\n\n");
        }
    } else {
        dr_fprintf(pt->output_file, "***No Fault 4***\n\n");
    }
}

// client want to do
inline void
DoWhatClientWantTodo(void *drcontext, per_thread_t *pt, context_handle_t cur_ctxt_hndl,
                     app_pc cur_addr, char type, ushort size)
{
  //dr_printf("cur addr is %x\n", (uint64_t)cur_addr);
   /*if (cur_addr == NULL) {
       (pt->access_NULL_Fault_1)->push_back(cur_ctxt_hndl);
       return ;
   }*/

   data_handle_t data_hndl = drcctlib_get_data_hndl(drcontext, cur_addr);

   if (data_hndl.object_type == UNKNOWN_OBJECT) {
       //(pt->access_unalloc_addr_Fault_2)->push_back(cur_ctxt_hndl);
       return ;
   }

   if (data_hndl.object_type == STACK_OBJECT || data_hndl.object_type == STATIC_OBJECT) {
       //do not deal with stack object and static object
       return ;
   }
   
   if (data_hndl.alloc == 0) {
       if (drcctlib_test_if_need_print_backtrace(cur_ctxt_hndl, 1) && (pt->access_unalloc_addr_Fault_2)->size() < 1000) {
            (pt->access_unalloc_addr_Fault_2)->push_back(cur_ctxt_hndl);   
       }
       return ;
   }

   if (data_hndl.freed == 1) {
       if (drcctlib_test_if_need_print_backtrace(cur_ctxt_hndl, 1) && (pt->access_freed_addr_Fault_3)->size() < 1000) {
           (pt->access_freed_addr_Fault_3)->push_back(cur_ctxt_hndl);
       }
       return ;
   }

    if (type == 0) {
        //read op
        if (data_hndl.init == 0 
                && drcctlib_test_if_need_print_backtrace(cur_ctxt_hndl, 1) && (pt->access_uninit_addr_Fault_4)->size() < 1000) {
            (pt->access_uninit_addr_Fault_4)->push_back(cur_ctxt_hndl);
        }
    } else {
        //write op
        if (data_hndl.init == 0) {
            data_handle_t *change_init_flag = mem_check_get_shadow_memory(cur_addr);
            for (int i = 0; i < size; i++) {
                change_init_flag[i].init = 1;
            }
        }
    }
}

static inline void
InstrumentPerBBCache(void *drcontext, context_handle_t ctxt_hndl, int32_t slot_num,
                     int32_t mem_ref_num, mem_ref_msg_t *mem_ref_start, void **data)
{
    per_thread_t *pt;
    if (*data != NULL) {
        pt = (per_thread_t *)*data;
    } else {
        pt = (per_thread_t *)drmgr_get_tls_field(drcontext, tls_idx);
        *data = pt;
    }

    for (int32_t i = 0; i < mem_ref_num; i++) {
        if (mem_ref_start[i].slot >= slot_num) {
            break;
        }
        DoWhatClientWantTodo(drcontext, pt, ctxt_hndl + mem_ref_start[i].slot,
                             mem_ref_start[i].addr, mem_ref_start[i].type, mem_ref_start[i].size);
    }
}

#ifdef DEBUG_REUSE
static void
ThreadDebugFileInit(per_thread_t *pt)
{
    int32_t id = drcctlib_get_thread_id();
    char name[MAXIMUM_FILEPATH] = "";
    sprintf(name + strlen(name), "%s/thread-%d.debug.log", g_folder_name.c_str(), id);
    pt->log_file = dr_open_file(name, DR_FILE_WRITE_APPEND | DR_FILE_ALLOW_LARGE);
    DR_ASSERT(pt->log_file != INVALID_FILE);
}
#endif

static void
ThreadOutputFileInit(per_thread_t *pt)
{
    int32_t id = drcctlib_get_thread_id();
    char name[MAXIMUM_FILEPATH] = "";
    sprintf(name + strlen(name), "%s/thread-%d.topn.log", g_folder_name.c_str(), id);
    pt->output_file = dr_open_file(name, DR_FILE_WRITE_OVERWRITE | DR_FILE_ALLOW_LARGE);
    DR_ASSERT(pt->output_file != INVALID_FILE);
}

static void
ClientThreadStart(void *drcontext)
{
    per_thread_t *pt = (per_thread_t *)dr_thread_alloc(drcontext, sizeof(per_thread_t));
    if (pt == NULL) {
        DRCCTLIB_EXIT_PROCESS("pt == NULL");
    }
    drmgr_set_tls_field(drcontext, tls_idx, (void *)pt);

    pt->access_NULL_Fault_1 = new vector<context_handle_t>();
    pt->access_unalloc_addr_Fault_2 = new vector<context_handle_t>();
    pt->access_freed_addr_Fault_3 = new vector<context_handle_t>();
    pt->access_uninit_addr_Fault_4 = new vector<context_handle_t>();
    pt->pt_id = drcctlib_get_thread_id();

    ThreadOutputFileInit(pt);
#ifdef DEBUG_REUSE
    ThreadDebugFileInit(pt);
#endif
}

static void
ClientThreadEnd(void *drcontext)
{
    per_thread_t *pt = (per_thread_t *)drmgr_get_tls_field(drcontext, tls_idx);
    PrintResult(pt);
    dr_close_file(pt->output_file);
    delete pt->access_NULL_Fault_1;
    delete pt->access_unalloc_addr_Fault_2;
    delete pt->access_freed_addr_Fault_3;
    delete pt->access_uninit_addr_Fault_4;
#ifdef DEBUG_REUSE
    dr_close_file(pt->log_file);
#endif
    dr_thread_free(drcontext, pt, sizeof(per_thread_t));
}

static void
ClientInit(int argc, const char *argv[])
{
    char name[MAXIMUM_FILEPATH] = "";
    DRCCTLIB_INIT_LOG_FILE_NAME(
        name, "drcctlib_mem_check", "out");
    g_folder_name.assign(name, strlen(name));
    mkdir(g_folder_name.c_str(), S_IRWXU | S_IRWXG | S_IROTH | S_IXOTH);
}

static void
ClientExit(void)
{
    drcctlib_check_memory_leak();
    drcctlib_exit();
    if (!drmgr_unregister_thread_init_event(ClientThreadStart) ||
        !drmgr_unregister_thread_exit_event(ClientThreadEnd) ||
        !drmgr_unregister_tls_field(tls_idx)) {
        DRCCTLIB_PRINTF(
            "ERROR: drcctlib_mem_check failed to unregister in ClientExit");
    }
    drmgr_exit();
}

#ifdef __cplusplus
extern "C" {
#endif

DR_EXPORT void
dr_client_main(client_id_t id, int argc, const char *argv[])
{
    dr_set_client_name("DynamoRIO Client 'drcctlib_mem_check'",
                       "http://dynamorio.org/issues");
    ClientInit(argc, argv);

    if (!drmgr_init()) {
        DRCCTLIB_EXIT_PROCESS(
            "ERROR: drcctlib_mem_check unable to initialize drmgr");
    }
    drmgr_priority_t thread_init_pri = { sizeof(thread_init_pri),
                                         "drcctlib_mem_check-thread_init", NULL, NULL,
                                         DRCCTLIB_THREAD_EVENT_PRI + 1 };
    drmgr_priority_t thread_exit_pri = { sizeof(thread_exit_pri),
                                         "drcctlib_mem_check-thread-exit", NULL, NULL,
                                         DRCCTLIB_THREAD_EVENT_PRI + 1 };
    drmgr_register_thread_init_event_ex(ClientThreadStart, &thread_init_pri);
    drmgr_register_thread_exit_event_ex(ClientThreadEnd, &thread_exit_pri);
    tls_idx = drmgr_register_tls_field();
    if (tls_idx == -1) {
        DRCCTLIB_EXIT_PROCESS(
            "ERROR: drcctlib_mem_check drmgr_register_tls_field fail");
    }
    drcctlib_init_ex(DRCCTLIB_FILTER_MEM_ACCESS_INSTR, INVALID_FILE, NULL, NULL,
                     InstrumentPerBBCache,
                     DRCCTLIB_COLLECT_DATA_CENTRIC_MESSAGE | DRCCTLIB_CACHE_MODE |
                         DRCCTLIB_CACHE_MEMEORY_ACCESS_ADDR);
    dr_register_exit_event(ClientExit);
}

#ifdef __cplusplus
}
#endif